/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
#ifndef CHANNEL_MANAGER_H
#define CHANNEL_MANAGER_H

#include "hccl_api.h"
#include "hccl_types.h"
#include "transport_pub.h"
#include "hccl_common.h"
#include "hccl_mem_defs.h"
#include "transport_pub.h"
#include "aicpu_operator_pub.h"
#include "channel_param.h"
#include "manager_common.h"

#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <string>

namespace std {
    template <>
    struct hash<ChannelDesc> {
        size_t operator()(const ChannelDesc& desc) const {
            size_t hash = 0;
            // 仅区分remoteRank和protocol
            hash ^= std::hash<uint32_t>()(desc.remoteRank);
            hash ^= std::hash<int32_t>()(static_cast<int32_t>(desc.protocol));
            return hash;
        }
    };
}

namespace hccl {

struct ChannelDescEqual {
    bool operator()(const ChannelDesc& lcd, const ChannelDesc& rcd) const {
        return lcd.remoteRank == rcd.remoteRank && lcd.protocol == rcd.protocol;
    }
};

class ChannelManager {
public:
    ChannelManager();
    ~ChannelManager();
    HcclResult Init(aclrtBinHandle binHandle, u32 userRank, const ManagerCallbacks& callbacks);
    HcclResult SetChannelCallbacks(const ChannelManagerCallbacks& channelCallbacks);
    HcclResult ChannelCommCreate(const std::string &commId, const std::string &tag, CommEngine engine, 
        const ChannelDesc *channelDescList, uint32_t listNum, ChannelHandle *channelList);
    HcclResult ChannelCommGetNotifyNum(ChannelHandle channel, uint32_t *notifyNum);
    HcclResult ChannelCommDestroy(ChannelHandle *channelList, uint32_t channelNum);
    HcclResult ChannelCommGetHcclBuffer(ChannelHandle channel, CommBuffer *buffer);
    HcclResult ChannelCommGetRemoteMem(ChannelHandle channel, HcclMem **remoteMem, uint32_t *memNum);

private:
    template <typename T>
    HcclResult CopyVectorToDeviceMem(const u64 len, DeviceMem &dstDeviceMem, const std::vector<T> &srcVec);
    HcclResult AllocAndClearHostMem(u64 size, std::shared_ptr<HostMem> &buffer) const;
    HcclResult CreateWorkSpace(u64 size, DeviceMem &buffer) const;
    HcclResult CheckNotifyOrQPMaxNum(u64 &existNum, const u64 &MaxNum, const bool &isNotifyRes);
    HcclResult DeepCopyH2DChannelP2p(const HcclChannelP2p &hostChannelP2p, HcclChannelP2p &deviceChannelP2p);
    HcclResult DeepCopyH2DChannelRoce(const HcclChannelRoce &hostChannelRoce, HcclChannelRoce &deviceChannelRoce);
    HcclResult DeepCopyH2DChannelRemoteResV2(const HcclIndOpChannelRemoteResV2 &hostRemoteResV2, 
        HcclIndOpChannelRemoteResV2 &deviceRemoteResV2);
    HcclResult DeepCopyH2DchannelParam(const HcclIndOpChannelRemoteResV3 &hostChannelParam, 
        HcclIndOpChannelRemoteResV3 &deviceChannelParam);
    HcclResult ReleaseChannelParam(HcclIndOpChannelRemoteResV3 &channelParam);
    HcclResult BuildOpRemoteChannelP2pResParam(const LINK &link, HcclIndOpChannelRemoteResV2 &remoteRes);
    HcclResult BuildOpRemoteChannelRoceResParam(const LINK &link, HcclIndOpChannelRemoteResV2 &remoteRes);
    HcclResult ParseChannelRemoteDataToMem(const OpCommTransport &opTransportResponse, HcclIndOpChannelRemoteResV3 &channelParam);
    HcclResult AicpuChannelInit(const std::string &commId, const std::string &tag, CommEngine engine, 
        const OpCommTransport &opTransportResponse, ChannelHandle *channelList, uint32_t listNum);
    void ClearOpTransportResponseLinks(OpCommTransport &opTransportResponse);
    OpCommTransport BuildChannelRequests(const std::vector<ChannelDesc> &descs);

    HcclResult CheckChannelParam(const std::string &tag, CommEngine engine, const ChannelDesc *channelDesc,
        uint32_t descNum);
    HcclResult RegisterHandle(const std::string& key, CommEngine engine, const ChannelDesc& channelDesc, ChannelHandle channelHandle);
    HcclResult RegisterHandleHDPair(ChannelHandle deviceChannelHandle, ChannelHandle hostChannelHandle);
    HcclResult UnregisterHandle(ChannelHandle channelHandle);
    HcclResult PrepareHandleArray(const std::string &tag, CommEngine engine, const ChannelDesc *channelDesc, 
        uint32_t descNum, ChannelHandle *channelHandleArray, std::vector<ChannelDesc> &needCreateDescs,
        std::vector<uint32_t> &needCreateIndices);
    HcclResult IsChannelExist(ChannelHandle channel);
    HcclResult GetHostChannel(ChannelHandle channel, ChannelHandle &hostChannel);
    
    std::unordered_map<std::string, ChannelHandle> channelHandleMap_;
    std::unordered_map<ChannelHandle, std::string> keyMap_;
    std::unordered_map<ChannelHandle, CommEngine> engineMap_;
    std::unordered_map<ChannelHandle, ChannelHandle> channelD2HMap_;
    std::unordered_set<ChannelHandle*> channelHandleArraySet_;
    std::unordered_map<ChannelHandle, LINK> linkMap_;
    std::vector<LINK> channelLinks_{};
    u32 userRank_;
    std::vector<RankInfo> rankInfoList_;
    std::vector<std::shared_ptr<DeviceMem>> channelParamMemVector_;
    std::list<DeviceMem> channelParamMemList_;
    aclrtBinHandle binHandle_;
    ManagerCallbacks callbacks_;  // 存储回调函数
    ChannelManagerCallbacks channelCallbacks_;  // channelMgr的回调函数
};

} // namespace hccl

#endif  // CHANNEL_MANAGER_H
